#!/bin/bash

# K8s服务器准备脚本
# 自动化配置Ubuntu Server系统环境

set -e

# 配置文件路径
CONFIG_FILE="./k8s-config.conf"
LOG_FILE="/var/log/k8s-prepare.log"

# 颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

# 日志函数
log() {
    local level=$1
    shift
    local message="$*"
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    
    case $level in
        "INFO")
            echo -e "${BLUE}[INFO]${NC} $message" | tee -a $LOG_FILE
            ;;
        "SUCCESS")
            echo -e "${GREEN}[SUCCESS]${NC} $message" | tee -a $LOG_FILE
            ;;
        "WARNING")
            echo -e "${YELLOW}[WARNING]${NC} $message" | tee -a $LOG_FILE
            ;;
        "ERROR")
            echo -e "${RED}[ERROR]${NC} $message" | tee -a $LOG_FILE
            ;;
    esac
}

# 错误处理函数
error_exit() {
    log "ERROR" "$1"
    log "ERROR" "服务器准备失败，请检查日志: $LOG_FILE"
    exit 1
}

# 检查配置文件
check_config() {
    if [ ! -f "$CONFIG_FILE" ]; then
        error_exit "配置文件不存在: $CONFIG_FILE"
    fi
    
    # 加载配置文件
    source "$CONFIG_FILE"
    
    log "SUCCESS" "配置文件加载成功"
    log "INFO" "集群名称: $CLUSTER_NAME"
    log "INFO" "K8s版本: $K8S_VERSION"
}

# 检查root权限
check_root() {
    if [ "$EUID" -ne 0 ]; then
        error_exit "请使用root用户运行此脚本"
    fi
    log "SUCCESS" "权限检查通过"
}

# 检查操作系统
check_os() {
    log "INFO" "检查操作系统..."
    if [ -f /etc/os-release ]; then
        . /etc/os-release
        if [[ "$ID" == "ubuntu" ]]; then
            log "SUCCESS" "操作系统检查通过: $PRETTY_NAME"
        else
            log "WARNING" "当前系统不是Ubuntu，可能不兼容"
        fi
    else
        log "WARNING" "无法确定操作系统版本"
    fi
}

# 更新系统包
update_system() {
    log "INFO" "更新系统包..."
    
    # 更新包列表
    apt-get update -y
    
    # 升级系统包
    apt-get upgrade -y
    
    # 安装必要工具
    apt-get install -y \
        curl \
        wget \
        gnupg \
        lsb-release \
        ca-certificates \
        apt-transport-https \
        software-properties-common \
        net-tools \
        vim \
        htop \
        iotop \
        ntp \
        chrony
    
    log "SUCCESS" "系统包更新完成"
}

# 配置时间同步
configure_ntp() {
    log "INFO" "配置时间同步..."
    
    # 安装并启动chrony
    systemctl enable chrony
    systemctl start chrony
    
    # 同步时间
    chrony sources -v
    chrony tracking
    
    log "SUCCESS" "时间同步配置完成"
}

# 配置主机名
configure_hostname() {
    local hostname=$1
    log "INFO" "配置主机名: $hostname"
    
    # 设置主机名
    hostnamectl set-hostname "$hostname"
    
    # 更新/etc/hosts
    local current_ip=$(hostname -I | awk '{print $1}')
    echo "$current_ip $hostname" >> /etc/hosts
    
    log "SUCCESS" "主机名配置完成: $hostname"
}

# 配置网络
configure_network() {
    log "INFO" "配置网络..."
    
    # 获取网络接口名称
    local interface=$(ip route | grep default | awk '{print $5}' | head -1)
    
    if [ -z "$interface" ]; then
        error_exit "无法确定网络接口"
    fi
    
    log "INFO" "网络接口: $interface"
    
    # 配置静态IP（如果需要）
    if [ -n "$MASTER_IP" ] && [ "$MASTER_IP" != "dhcp" ]; then
        log "INFO" "配置静态IP: $MASTER_IP"
        
        # 备份原始配置
        cp /etc/netplan/*.yaml /etc/netplan/*.yaml.backup 2>/dev/null || true
        
        # 创建网络配置
        cat > /etc/netplan/01-netcfg.yaml << EOF
network:
  version: 2
  renderer: networkd
  ethernets:
    $interface:
      dhcp4: false
      addresses:
        - $MASTER_IP/24
      gateway4: $GATEWAY
      nameservers:
        addresses: [$DNS_SERVERS]
EOF
        
        # 应用配置
        netplan apply
    fi
    
    log "SUCCESS" "网络配置完成"
}

# 配置防火墙
configure_firewall() {
    log "INFO" "配置防火墙..."
    
    if [ "$FIREWALL_ENABLED" = "false" ]; then
        # 禁用防火墙
        ufw --force disable
        log "SUCCESS" "防火墙已禁用"
    else
        # 配置防火墙规则
        ufw --force enable
        ufw default deny incoming
        ufw default allow outgoing
        
        # 允许SSH
        ufw allow ssh
        
        # 允许K8s端口
        ufw allow 6443/tcp  # Kubernetes API server
        ufw allow 2379:2380/tcp  # etcd server client API
        ufw allow 10250/tcp  # Kubelet API
        ufw allow 10251/tcp  # kube-scheduler
        ufw allow 10252/tcp  # kube-controller-manager
        ufw allow 30000:32767/tcp  # NodePort Services
        
        log "SUCCESS" "防火墙规则配置完成"
    fi
}

# 配置swap
configure_swap() {
    log "INFO" "配置swap..."
    
    if [ "$SWAP_ENABLED" = "false" ]; then
        # 关闭swap
        swapoff -a
        sed -i '/ swap / s/^\(.*\)$/#\1/g' /etc/fstab
        log "SUCCESS" "swap已关闭"
    else
        log "INFO" "保持swap开启状态"
    fi
}

# 配置内核参数
configure_kernel() {
    log "INFO" "配置内核参数..."
    
    # 创建内核参数配置
    cat > /etc/sysctl.d/k8s.conf << EOF
# Kubernetes内核参数
net.bridge.bridge-nf-call-ip6tables = 1
net.bridge.bridge-nf-call-iptables = 1
net.ipv4.ip_forward = 1
vm.swappiness = 0
vm.overcommit_memory = 1
vm.panic_on_oom = 0
net.ipv4.tcp_keepalive_time = 600
net.ipv4.tcp_keepalive_intvl = 60
net.ipv4.tcp_keepalive_probes = 10
net.ipv6.conf.all.disable_ipv6 = 1
net.ipv6.conf.default.disable_ipv6 = 1
net.ipv6.conf.lo.disable_ipv6 = 1
EOF
    
    # 应用配置
    sysctl --system
    
    # 加载必要模块
    modprobe br_netfilter
    echo 'br_netfilter' > /etc/modules-load.d/k8s.conf
    
    # 加载overlay模块
    modprobe overlay
    echo 'overlay' >> /etc/modules-load.d/k8s.conf
    
    log "SUCCESS" "内核参数配置完成"
}

# 安装Docker
install_docker() {
    log "INFO" "安装Docker..."
    
    if command -v docker &> /dev/null; then
        log "INFO" "Docker已安装，跳过安装步骤"
        return
    fi
    
    # 添加Docker官方GPG密钥
    curl -fsSL https://download.docker.com/linux/ubuntu/gpg | gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg
    
    # 添加Docker仓库
    echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/ubuntu $(lsb_release -cs) stable" | tee /etc/apt/sources.list.d/docker.list > /dev/null
    
    # 更新包列表
    apt-get update
    
    # 安装Docker
    apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
    
    # 启动Docker服务
    systemctl enable docker
    systemctl start docker
    
    # 配置Docker
    mkdir -p /etc/docker
    cat > /etc/docker/daemon.json << EOF
{
  "exec-opts": ["native.cgroupdriver=systemd"],
  "log-driver": "json-file",
  "log-opts": {
    "max-size": "100m"
  },
  "storage-driver": "overlay2"
}
EOF
    
    # 重启Docker服务
    systemctl daemon-reload
    systemctl restart docker
    
    log "SUCCESS" "Docker安装完成"
}

# 配置SSH
configure_ssh() {
    log "INFO" "配置SSH..."
    
    # 备份SSH配置
    cp /etc/ssh/sshd_config /etc/ssh/sshd_config.backup
    
    # 配置SSH
    sed -i 's/#PermitRootLogin prohibit-password/PermitRootLogin yes/' /etc/ssh/sshd_config
    sed -i 's/#PasswordAuthentication yes/PasswordAuthentication yes/' /etc/ssh/sshd_config
    
    # 重启SSH服务
    systemctl restart ssh
    
    log "SUCCESS" "SSH配置完成"
}

# 创建用户和目录
create_directories() {
    log "INFO" "创建必要目录..."
    
    # 创建K8s相关目录
    mkdir -p /etc/kubernetes
    mkdir -p /var/lib/kubelet
    mkdir -p /var/lib/kube-proxy
    mkdir -p /var/lib/etcd
    mkdir -p /opt/cni/bin
    mkdir -p /etc/cni/net.d
    
    # 设置权限
    chmod 755 /etc/kubernetes
    chmod 755 /var/lib/kubelet
    chmod 755 /var/lib/kube-proxy
    chmod 755 /var/lib/etcd
    chmod 755 /opt/cni/bin
    chmod 755 /etc/cni/net.d
    
    log "SUCCESS" "目录创建完成"
}

# 配置日志轮转
configure_logrotate() {
    log "INFO" "配置日志轮转..."
    
    cat > /etc/logrotate.d/k8s << EOF
/var/log/k8s-*.log {
    daily
    missingok
    rotate $LOG_RETENTION_DAYS
    compress
    delaycompress
    notifempty
    create 644 root root
}
EOF
    
    log "SUCCESS" "日志轮转配置完成"
}

# 生成系统信息报告
generate_report() {
    local report_file="/tmp/k8s-prepare-report-$(date +%Y%m%d-%H%M%S).txt"
    
    log "INFO" "生成系统准备报告: $report_file"
    
    {
        echo "K8s服务器准备报告"
        echo "生成时间: $(date)"
        echo "主机名: $(hostname)"
        echo "操作系统: $(cat /etc/os-release | grep PRETTY_NAME | cut -d'=' -f2 | tr -d '\"')"
        echo "=========================================="
        echo ""
        
        echo "系统信息:"
        echo "内核版本: $(uname -r)"
        echo "架构: $(uname -m)"
        echo "内存: $(free -h | grep Mem | awk '{print $2}')"
        echo "磁盘: $(df -h / | awk 'NR==2{print $2}')"
        echo ""
        
        echo "网络配置:"
        ip addr show
        echo ""
        
        echo "服务状态:"
        systemctl status docker chrony --no-pager
        echo ""
        
        echo "内核参数:"
        sysctl -a | grep -E "(bridge|ip_forward|swappiness)" | head -10
        echo ""
        
    } > $report_file
    
    log "SUCCESS" "系统准备报告已生成: $report_file"
}

# 显示使用说明
show_usage() {
    echo "K8s服务器准备脚本"
    echo ""
    echo "用法:"
    echo "  $0 [选项]"
    echo ""
    echo "选项:"
    echo "  --hostname=NAME     设置主机名"
    echo "  --ip=IP             设置IP地址"
    echo "  --skip-docker       跳过Docker安装"
    echo "  --skip-network      跳过网络配置"
    echo "  --help              显示帮助信息"
    echo ""
    echo "示例:"
    echo "  $0 --hostname=master-node --ip=192.168.1.10"
    echo "  $0 --hostname=worker-node1 --ip=192.168.1.11"
}

# 主函数
main() {
    log "INFO" "开始K8s服务器准备..."
    log "INFO" "准备日志: $LOG_FILE"
    
    # 解析命令行参数
    local custom_hostname=""
    local custom_ip=""
    local skip_docker=false
    local skip_network=false
    
    while [[ $# -gt 0 ]]; do
        case $1 in
            --hostname=*)
                custom_hostname="${1#*=}"
                shift
                ;;
            --ip=*)
                custom_ip="${1#*=}"
                shift
                ;;
            --skip-docker)
                skip_docker=true
                shift
                ;;
            --skip-network)
                skip_network=true
                shift
                ;;
            --help)
                show_usage
                exit 0
                ;;
            *)
                echo "未知参数: $1"
                show_usage
                exit 1
                ;;
        esac
    done
    
    # 执行准备步骤
    check_root
    check_config
    check_os
    update_system
    configure_ntp
    
    # 配置主机名
    if [ -n "$custom_hostname" ]; then
        configure_hostname "$custom_hostname"
    fi
    
    # 配置网络
    if [ "$skip_network" = false ]; then
        configure_network
    fi
    
    configure_firewall
    configure_swap
    configure_kernel
    
    # 安装Docker
    if [ "$skip_docker" = false ]; then
        install_docker
    fi
    
    configure_ssh
    create_directories
    configure_logrotate
    
    # 生成报告
    generate_report
    
    log "SUCCESS" "K8s服务器准备完成！"
    log "INFO" "请重启服务器以确保所有配置生效"
}

# 运行主函数
main "$@"
